#!/usr/bin/python3
# -*- coding: utf-8 -*-
# function: prepare data
# author: changhengyi
# date: 2020.12.4

import os
import sys
import numpy as np
import time
import random
# from python_speech_features import logfbank
# import scipy.io.wavfile as wav
# import wave

cmu_phones = {'sil': 0, 'aa': 1, 'ae': 2, 'ah': 3, 'ao': 4, 'aw': 5, 'ay': 6, 'b': 7, 'ch': 8, 'd': 9, 'dh': 10, \
              'eh': 11, 'er': 12, 'ey': 13, 'f': 14, 'g': 15, 'hh': 16, 'ih': 17, 'iy': 18, 'jh': 19, 'k': 20, 'l': 21, \
              'm': 22, 'n': 23, 'ng': 24, 'ow': 25, 'oy': 26, 'p': 27, 'r': 28, 's': 29, 'sh': 30, 't': 31, 'th': 32, \
              'uh': 33, 'uw': 34, 'v': 35, 'w': 36, 'y': 37, 'z': 38, 'zh': 39}

map_60_39 = {'aa': 1, 'ae': 2, 'ah': 3, 'ao': 4, 'aw': 5, 'ax': 3, 'ax-h': 3, 'axr': 12, 'ay': 6, 'b': 7, \
             'bcl': 0, 'ch': 8, 'd': 9, 'dcl': 0, 'dh': 10, 'dx': 31, 'eh': 11, 'el': 21, 'em': 22, 'en': 23, \
             'eng': 24, 'epi': 0, 'er': 12, 'ey': 13, 'f': 14, 'g': 15, 'gcl': 0, 'h#': 0, 'hh': 16, 'hv': 16,\
             'ih': 17, 'ix': 17, 'iy': 18, 'jh': 19, 'k': 20, 'kcl': 0, 'l': 21, 'm': 22, 'n': 23, 'ng': 24, \
             'nx': 23, 'ow': 25, 'oy': 26, 'p': 27, 'pau': 0, 'pcl': 0, 'q': 0, 'r': 28, 's': 29, 'sh': 30, \
             't': 31, 'tcl': 0, 'th': 32, 'uh': 33, 'uw': 34, 'ux': 34, 'v': 35, 'w': 36, 'y': 37, 'z': 38, 'zh': 39}


class DataLoader():
    def __init__(self, train_path, test_path):
        self.train_path = train_path
        self.test_path = test_path
        self.train_data = []
        self.test_data = []
        self.index = 0
        
        with open(self.train_path, "r") as f:
            print("加载训练数据...")
            line = f.readline()
            while line:
                self.train_data.append(line.strip().split())
                line = f.readline()
        with open(self.test_path, "r") as f:
            print("加载测试数据...")
            line = f.readline()
            while line:
                self.test_data.append(line.strip().split())
                line = f.readline()
        random.shuffle(self.train_data)
        random.shuffle(self.test_data)
        print("Dataset shuffled....")
        

    def get_batch(self, batch_size, audio_len, feat_dim):
        total_len = batch_size * audio_len
        len_sum = 0
        voices = []
        labels = []
        while len_sum < total_len:
            data = self.train_data[self.index]
            self.index += 1
            if self.index == len(self.train_data):
                self.index = 0
                random.shuffle(self.train_data)
                print("Dataset shuffled....")
            voice = np.load(data[1])
            label = np.load(data[2])
            min_len = min(voice.shape[0], label.shape[0])
            voice = voice[:min_len]
            label = label[:min_len]
            len_sum += min_len

            voices.extend(voice)
            labels.extend(label)
        
        voices = np.array(voices[:total_len])
        voices.resize((batch_size, audio_len, feat_dim))
        labels = np.array(labels)[:total_len]
        labels.resize((batch_size, audio_len))

        return voices, labels


